import os
from zhipuai import ZhipuAI
from dotenv import load_dotenv, find_dotenv
import requests
import json
import streamlit as st

_ = load_dotenv(find_dotenv())
api_key = os.environ.get('ZHIPUAI_API_KEY')
amap_key = os.environ.get('amap_key')
if api_key is None:
    raise ValueError("API Key is not set in the .env file")
client = ZhipuAI(api_key=api_key)

def get_completion(messages, model="glm-4"):
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0.01,  
        tool_choice="auto", 
        tools=[{
            "type": "function",
            "function": {
                "name": "get_location_coordinate",
                "description": "根据地点关键字，获得地点的经纬度坐标",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "keywords": {
                            "type": "string",
                            "description": "地点关键字，需要被检索的地点文本信息，必须是中文",
                        },
                        "region": {
                            "type": "string",
                            "description": "POI所在的城市名或区域名，必须是中文，可选参数",
                        }
                    },
                    "required": ["keywords"],
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "search_nearby_pois",
                "description": "搜索给定坐标附近的POI",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "longitude": {
                            "type": "string",
                            "description": "中心点的经度",
                        },
                        "latitude": {
                            "type": "string",
                            "description": "中心点的纬度",
                        },
                        "keyword": {
                            "type": "string",
                            "description": "目标POI的关键字",
                        }
                    },
                    "required": ["longitude", "latitude", "keyword"],
                }
            }
        }],
    )
    return response.choices[0].message

def get_location_coordinate(keywords, region=""):
    url = f"""https://restapi.amap.com/v5/place/text?key={
        amap_key}&keywords={keywords}&region={region}"""
    ret = requests.get(url)
    result = ret.json()
    if "pois" in result and result["pois"]:
        return result["pois"]
    return None


def search_nearby_pois(longitude, latitude, keyword):
    url = f"""https://restapi.amap.com/v5/place/around?key={
        amap_key}&keywords={keyword}&location={longitude},{latitude}"""
    ret = requests.get(url)
    result = ret.json()
    ans = ""
    if "pois" in result and result["pois"]:
        for i in range(min(3, len(result["pois"]))):
            name = result["pois"][i]["name"]
            address = result["pois"][i]["address"]
            distance = result["pois"][i]["distance"]
            ans += f"{name}\n{address}\n距离：{distance}米\n\n"
    return ans


# 头像配置
ICON_AI = '💻'
ICON_USER = '🧑'

# 显示一条消息（包含头像与消息内容）
def dspMessage(role, content):
    with st.chat_message(role, avatar=ICON_AI if role == 'assistant' else ICON_USER):
        st.write(content)
        

# 追加并显示一条消息
def append_and_show(role, content):
    st.session_state.messages.append({"role": role, "content": content})      
    dspMessage(role, content)

if 'messages' not in st.session_state:
    st.session_state.messages = [{"role": "assistant", "content": "我是你的地图小助手，请问你查询什么地理信息？"}]


# 将会话中的messages列表中的消息全部显示出来
for msg in st.session_state.messages:
    dspMessage(msg["role"], msg["content"])

# 接受用户输入的提示词，并调用大模型API获得反馈
if prompt := st.chat_input():    
    append_and_show("user", prompt)

    messages = [
        {"role": "system", "content": """你是一个地图小助手，你应当先分析用户消息中的地址和关键字，然后调用API查找地址或关键字的经纬度，然后根据经纬度调用API查找用户消息中POI的信息。
         如果存在多个满足条件的地址及其对应经纬度，你必须要一次性将这些地址的经纬度和调用的API返回给应用查找对应的POI，不要遗漏。"""},
        {"role": "user", "content": prompt}
    ]

    response = get_completion(messages)
    print(f"\n首次调用结果：{response}\n")

    while (response.tool_calls is not None):
        messages.append(response.model_dump())  # 把大模型的回复加入到对话中
        for tool_call in response.tool_calls:
            args = json.loads(tool_call.function.arguments)
            print(f"\n调用API:{tool_call.function.name}参数：{args}")

            if (tool_call.function.name == "get_location_coordinate"):
                result = get_location_coordinate(**args)
            elif (tool_call.function.name == "search_nearby_pois"):
                result = search_nearby_pois(**args)

            if result is None or result == "":
                result = "未找到相应信息。" 
            print(f"调用API结果：{result}")

            messages.append({
                "tool_call_id": tool_call.id,  # 用于标识函数调用的 ID
                "role": "tool",
                "name": tool_call.function.name,
                "content": f"{result}" 
            })

        response = get_completion(messages)
        print(f"\n新一轮结果：{response}\n")

    # 保存并显示已经完成的回复
    append_and_show("assistant", response.content) 
    print(f"\n===END===\n")

